import torch
from torch import nn

from models import resnet


def get_module_name(name):
	name = name.split('.')
	if name[0] == 'module':
		i = 1
	else:
		i = 0

	if name[i] == 'features':
		i += 1

	return name[i]


def get_fine_tuning_parameters(model, ft_begin_module):
	if not ft_begin_module:
		return model.parameters()

	parameters = []

	add_flag = False
	for k, v in model.named_parameters():
		if ft_begin_module == get_module_name(k):
			add_flag = True

		if add_flag:
			parameters.append({'params': v})

	return parameters


def generate_model(opt):
	model = resnet.generate_model(model_depth=opt.model_depth,
	                              n_classes=opt.n_classes,
	                              shortcut_type=opt.resnet_shortcut,
	                              no_max_pool=opt.no_max_pool,
	                              widen_factor=opt.resnet_widen_factor)

	return model


def load_pretrained_model(model, pretrain_path, model_name, n_finetune_classes):
	if pretrain_path:
		print('loading pretrained model {}'.format(pretrain_path))
		pretrain = torch.load(pretrain_path, map_location='cpu')

		model.load_state_dict(pretrain['state_dict'])
		tmp_model = model

		if model_name == 'densenet':
			tmp_model.classifier = nn.Linear(tmp_model.classifier.in_features,
			                                 n_finetune_classes)

		else:
			tmp_model.fc = nn.Linear(tmp_model.fc.in_features, n_finetune_classes)

	return model


def make_data_parallel(model, is_distributed, device):
	if is_distributed:
		if device.type == 'cuda' and device.index is not None:

			torch.cuda.set_device(device)
			model.to(device)

			model = nn.parallel.DistributedDataParallel(model, device_ids=[device])

		else:
			model.to(device)
			model = nn.parallel.DistributedDataParallel(model)

	elif device.type == 'cuda':
		model = nn.DataParallel(model, device_ids=None).cuda()

	return model
